// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::{hash_map::Entry, HashMap, VecDeque}, sync::Arc, time::Duration};

use lazy_static::lazy_static;
use rosrust::{Publisher, RawMessage, RawMessageDescription, Subscriber};
use tokio::sync::{Mutex, Semaphore};

lazy_static! {
    static ref MSG_TYPES: Arc<Mutex<HashMap<String, String>>> = Arc::new(Mutex::new(HashMap::new()));
    static ref RUNTIME: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread()
        .build()
        .expect("Failed to build tokio runtime");
}

pub struct RosNode {
    publishers: HashMap<String, Publisher<RawMessage>>,
    subscribers: Arc<Mutex<HashMap<String, SubscriberData>>>,
}

struct SubscriberData {
    _subscriber: Subscriber,
    queue: VecDeque<Vec<u8>>,
    semaphore: Arc<Semaphore>,
}

#[derive(Debug, thiserror::Error)]
pub enum RosNodeError {
    #[error("Failed to initialize ROS node")]
    InitFailed(#[source] rosrust::error::Error),
    #[error("Failed to scan message types")]
    ScanMsgTypesFailed(#[source] rosrust::error::ResponseError),
    #[error("Failed to create publisher")]
    CreatePublisherFailed(#[source] rosrust::error::Error),
    #[error("Failed to create subscriber")]
    CreateSubscriberFailed(#[source] rosrust::error::Error),
    #[error("Failed to publish message")]
    PublishFailed(#[source] rosrust::error::Error),
    #[error("Failed to subscribe to message")]
    SubscribeFailed,
    #[error("Message type not found")]
    MsgTypeNotFound(String),
    #[error("Failed to build tokio runtime")]
    TokioRuntimeBuildFailed(#[source] tokio::io::Error),
}

impl RosNode {
    const DEFAULT_QUEUE_SIZE: usize = 10;

    pub fn new(node_path: &str) -> Result<RosNode, RosNodeError> {
        rosrust::try_init(node_path).map_err(RosNodeError::InitFailed)?;
        Ok(RosNode {
            publishers: HashMap::new(),
            subscribers: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    pub fn send(&mut self, path: &str, data: &[u8]) -> Result<(), RosNodeError> {
        RUNTIME.block_on(self.send_async(path, data))
    }

    pub fn recv(&mut self, path: &str, timeout: u32) -> Result<Vec<u8>, RosNodeError> {
        RUNTIME.block_on(self.recv_async(path, timeout))
    }

    pub fn on_recv<F>(&mut self, path: &str, callback: F) -> Result<(), RosNodeError>
    where
        F: Fn(&mut [u8]) + Send + 'static,
    {
        RUNTIME.block_on(self.on_recv_async(path, callback))
    }

    pub async fn send_async(&mut self, path: &str, data: &[u8]) -> Result<(), RosNodeError> {
        let msg_type = Self::get_or_scan_msg_type(path).await?;
        let publisher = self.get_or_create_publisher(path, msg_type)?;
        publisher.send(RawMessage(data.to_vec())).map_err(RosNodeError::PublishFailed)?;
        Ok(())
    }

    pub async fn recv_async(&mut self, path: &str, timeout: u32) -> Result<Vec<u8>, RosNodeError> {
        let path_clone = path.to_string();
        let subscribers_clone = self.subscribers.clone();
        let semaphore = self.get_or_create_subscriber(
            path,
            move |raw_msg: RawMessage| {
                RUNTIME.block_on(async {
                    let mut subscribers = subscribers_clone.lock().await;
                    if let Some(subscriber) = subscribers.get_mut(&path_clone) {
                        subscriber.queue.push_back(raw_msg.0);
                        subscriber.semaphore.add_permits(1);
                    }
                });
            }
        ).await?;

        tokio::select! {
            _ = semaphore.acquire() => {
                let data = self.subscribers.lock()
                    .await
                    .get_mut(path)
                    .unwrap()
                    .queue
                    .pop_front()
                    .ok_or(RosNodeError::SubscribeFailed)?;
                Ok(data)
            }
            _ = tokio::time::sleep(Duration::from_millis(timeout as u64)) => {
                Err(RosNodeError::SubscribeFailed)
            }
        }
    }

    pub async fn on_recv_async<F>(&mut self, path: &str, callback: F) -> Result<(), RosNodeError>
    where
        F: Fn(&mut [u8]) + Send + 'static,
    {
        self.subscribers.lock().await.insert(
            path.to_string(),
            SubscriberData {
                _subscriber: Self::create_subscriber(path, move |mut raw_msg| {
                    let data = raw_msg.0.as_mut_slice();
                    callback(data);
                })?,
                queue: VecDeque::new(),
                semaphore: Arc::new(Semaphore::new(0)),
            },
        );
        Ok(())
    }

    fn get_or_create_publisher(&mut self, path: &str, msg_type: String) -> Result<&mut Publisher<RawMessage>, RosNodeError> {
        match self.publishers.entry(path.to_string()) {
            Entry::Occupied(entry) => Ok(entry.into_mut()),
            Entry::Vacant(entry) => {
                let publisher = Self::create_publisher(path, msg_type)?;
                Ok(entry.insert(publisher))
            }
        }
    }

    async fn get_or_create_subscriber<F>(
        &mut self,
        path: &str,
        callback: F,
    ) -> Result<Arc<Semaphore>, RosNodeError>
    where
        F: Fn(RawMessage) + Send + 'static,
    {
        match self.subscribers.lock().await.entry(path.to_string()) {
            Entry::Occupied(entry) => Ok(entry.get().semaphore.clone()),
            Entry::Vacant(entry) => {
                let subscriber = Self::create_subscriber(path, callback)?;
                let semaphore = Arc::new(Semaphore::new(0));
                entry.insert(SubscriberData {
                    _subscriber: subscriber,
                    queue: VecDeque::new(),
                    semaphore: semaphore.clone(),
                });
                Ok(semaphore)
            }
        }
    }

    fn create_publisher(topic: &str, msg_type: String) -> Result<Publisher<RawMessage>, RosNodeError> {
        rosrust::publish_with_description::<RawMessage>(
            topic,
            Self::DEFAULT_QUEUE_SIZE,
            RawMessageDescription {
                msg_definition: "*".to_string(),
                md5sum: "*".to_string(),
                msg_type,
            },
        )
        .map_err(RosNodeError::CreatePublisherFailed)
    }

    fn create_subscriber<F>(topic: &str, callback: F) -> Result<Subscriber, RosNodeError>
    where
        F: Fn(RawMessage) + Send + 'static,
    {
        rosrust::subscribe(
            topic,
            Self::DEFAULT_QUEUE_SIZE,
            callback,
        )
        .map_err(RosNodeError::CreateSubscriberFailed)
    }

    async fn get_or_scan_msg_type(path: &str) -> Result<String, RosNodeError> {
        let msg_type = match MSG_TYPES.lock().await.get(path) {
            Some(msg_type) => msg_type.to_string(),
            None => {
                let msg_types = Self::scan_msg_types()?;
                MSG_TYPES.lock().await.extend(msg_types);
                MSG_TYPES.lock().await.get(path)
                    .ok_or(RosNodeError::MsgTypeNotFound(path.to_string()))?
                    .to_string()
            },
        };
        Ok(msg_type)
    }

    fn scan_msg_types() -> Result<HashMap<String, String>, RosNodeError> {
        let topics = rosrust::topics().map_err(RosNodeError::ScanMsgTypesFailed)?;
        Ok(topics.into_iter().map(|topic| (topic.name, topic.datatype)).collect())
    }
}
